import glob
import os
import cv2
import numpy as np
import random
from PyUIC.AGI_TestWidget import Ui_Form
from AGI_Widgets.QPaintWidgetView import QPaintWidgetView
from PyQt5.QtWidgets import QWidget,QTableWidgetItem,QFileDialog
from PyQt5.QtCore import QThread
from PyQt5.Qt import pyqtSignal

from AI.AI_Classify import AI_Classify
# from AI.AI_Det import AI_Det
from AI_RC.RC_classify import C_Classify,R_Classify,All_Classify,D_Classify,Q_Classify
from AI_ICCN.ICCN_Det import ICCN_Det

from utils.File_reader import write_xml
class ReadThread(QThread):
    trigged = pyqtSignal(str)
    image_sig = pyqtSignal(list)
    cls_sig = pyqtSignal()
    def __init__(self,img_folder):
        super().__init__()

        self.img_folder = img_folder
        self.img_paths = []

    def run(self):
        jpg_paths = glob.glob(os.path.join(self.img_folder, "*.jpg"))
        self.img_paths.extend(jpg_paths)
        png_paths = glob.glob(os.path.join(self.img_folder, "*.png"))
        self.img_paths.extend(png_paths)
        bmp_paths = glob.glob(os.path.join(self.img_folder, "*.bmp"))
        self.img_paths.extend(bmp_paths)
        # print(self.img_paths)
        for i,x in enumerate(self.img_paths):
            self.trigged.emit(x)
            try:

                if i == 0:
                    img = cv2.imdecode(np.fromfile(x, dtype=np.uint8), -1)
                    self.image_sig.emit([img])
            except Exception as e:
                print(e)
            # self.cls_sig.emit()

class AGITestWidget(QWidget,Ui_Form):
    def __init__(self,parent=None):
        super(AGITestWidget,self).__init__(parent)
        self.setupUi(self)
        self.PaintWidgetCLS = QPaintWidgetView()
        self.PaintWidgetDET = QPaintWidgetView()
        self.gridLayout_cls.addWidget(self.PaintWidgetCLS)
        self.gridLayout_det.addWidget(self.PaintWidgetDET)

        ######param
        self.is_cls = False
        self.is_det = False
        # [classify_results,yolo_results]
        self.results = {'cls':[],'res':[]}
        self.img_path = ''
        self.img_folder = ''
        self.save_folder = ''
        self.img_paths = []
        self.cls_res = ''
        self.image = None
        self.T_SaveButton.hide()
        self.T_AutoSave_checkBox.hide()

        #信号槽
        self.connection_slot()

        #AI模型
        # self.AI_Classify=AI_Classify()
        # self.AI_Yolo = AI_Det()
        self.com_classity = All_Classify()
        self.c_classify = C_Classify()
        self.r_classify = R_Classify()
        self.d_classify = D_Classify()
        self.q_classify = Q_Classify()
        self.iccn_det = ICCN_Det()

    def connection_slot(self):
        self.T_OpenImag_Button.clicked.connect(self.OpenImage)
        self.T_OpenFolder_Button.clicked.connect(self.OpenFolder)
        self.T_SaveFolder_Button.clicked.connect(self.SaveFolder)
        self.T_SaveButton.clicked.connect(self.Save_Xml)

        # self.cls_checkBox.toggled.connect(self.Start_Classify)
        # self.det_checkBox.toggled.connect(self.Start_Yolo)

        # self.img_listWidget.itemClicked.connect(self.Select_Image)
        self.img_listWidget.currentRowChanged.connect(self.Row_Changed)
        self.T_AutoRunButton.clicked.connect(self.Auto_Run_AI)


    def OpenImage(self):
        self.clear_listwidget()
        self.img_path,img_type = QFileDialog.getOpenFileName(self,"选择图片",self.img_folder,"*.jpg;;*.png;;*.bmp;;All Files(*)")
        basename = os.path.basename(self.img_path)
        if self.img_path=='':
            return
        self.img_listWidget.addItem(self.img_path)
        self.image = cv2.imdecode(np.fromfile(self.img_path, dtype=np.uint8), -1)
        self.show_img([self.image])
        self.Start_Classify()
        # self.Start_Yolo()
        # self.Auto_Save_Xml()
        pass

    def OpenFolder(self):
        self.clear_listwidget()
        self.img_listWidget.currentRowChanged.connect(self.Row_Changed)
        self.img_folder = QFileDialog.getExistingDirectory(self, "选择文件夹", self.img_folder)
        self.read_img = ReadThread(self.img_folder)
        self.read_img.start()
        # self.read_img.wait()
        self.read_img.trigged.connect(self.add_listwidget)
        self.read_img.image_sig.connect(self.show_img)


        # jpg_paths = glob.glob(os.path.join(self.img_folder,"*.jpg"))
        # self.img_paths.extend(jpg_paths)
        # png_paths = glob.glob(os.path.join(self.img_folder,"*.png"))
        # self.img_paths.extend(png_paths)
        # bmp_paths = glob.glob(os.path.join(self.img_folder,"*.bmp"))
        # self.img_paths.extend(bmp_paths)
        # for x in self.img_paths:
        #     self.img_listWidget.addItem(x)
        #     self.show_img(x)
        #     self.Start_Classify()



    def SaveFolder(self):
        self.save_folder = QFileDialog.getExistingDirectory(self, "选择文件夹", self.save_folder)
        pass

    def show_img(self,image):

        self.PaintWidgetCLS.open_image(image[0])
        self.PaintWidgetDET.open_image(image[0])
        pass

    def Start_Classify(self):
        # if self.cls_checkBox.isChecked():
        if self.image is None:
            return
        if len(self.image.shape) !=3:
            return
        result1,conf1 = self.com_classity.run_img(self.image)
        result2 = ''
        conf2 = ''
        if result1 == 'IC' or result1 == 'CN':
            image,reses = self.iccn_det.det_image(self.image)
            image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
            self.PaintWidgetDET.open_image(image)
            text = ''
            for res in reses:
                temp = "{},{:.2f},x1:{:.2f},y1{:.2f},x2:{:.2f},y2:{:.2f}\n".format(
                    res[0],res[1],res[2][0],res[2][1],res[2][2],res[2][3])
                text = text+temp
            self.det_result_label.setText(text)
            self.results['res'] = reses
            if text == '':
                result2='OK'
            else:
                result2='NG'

        elif result1 == "R":
            result2,conf2 = self.r_classify.run_img(self.image)
            if result2 in ['R105','R1','R2','R3','R4','R5','R6']:
                result2 = 'OK'
            else:
                result2='NG'
        elif result1 == 'C':
            result2,conf2 = self.c_classify.run_img(self.image)
            if result2 in ['1','2','3','4','5','6','7','8','9','10','11','12']:
                result2 = 'OK'
            else:
                result2='NG'
        elif result1 == "D":
            result2, conf2 = self.d_classify.run_img(self.image)
            if result2 in ['D1','D2','D3','D4','D5']:
                result2 = 'OK'
            else:
                result2 = 'NG'
        elif result1 == "Q":
            result2, conf2 = self.q_classify.run_img(self.image)
            if result2 in ['Q1', 'Q2', 'Q3', 'Q4']:
                result2 = 'OK'
            else:
                result2 = 'NG'

            # result1 =random.sample(['CN1','CN2','IC1','IC2'],1)[0]
            # conf1 = 0.9
            # if result1 =='CN2' or result1 =='IC2':
            #     result2 = 'OK'
            #     conf2 = 1.0
            # else:
            #     result2,conf2 = self.AI_Classify.classify_NG(self.image,result1)
            # result2 = random.sample(['NG1','NG2','NG3','NG4'],1)[0]
            # conf2 = 0.9
            # self.results['cls'] = [result1,result2]
        self.results['cls'] = [result1,result2]
        self.cls_result_label.setText('class:{},conf:{}\nOK/NG:{},conf:{}'.format(result1,conf1,result2,conf2))

    # def Start_Yolo(self):
    #     if self.det_checkBox.isChecked():
    #         if self.image is None:
    #             return
    #         if len(self.image.shape) !=3:
    #             return
    #         image,reses = self.AI_Yolo.run_image(self.image)
    #         image = cv2.cvtColor(np.asarray(image),cv2.COLOR_RGB2BGR)
    #         self.PaintWidgetDET.open_image(image)
    #         text = ''
    #         for res in reses:
    #             temp = "{},{:.2f},x1:{:.2f},y1{:.2f},x2:{:.2f},y2:{:.2f}\n".format(
    #                 res[0],res[1],res[2][0],res[2][1],res[2][2],res[2][3])
    #             text = text+temp
    #         self.det_result_label.setText(text)
    #         self.results['res'] = reses
    #         result = ['NG1',[20,20,50,50],0.8]

    def add_listwidget(self,x):
        self.img_listWidget.addItem(x)

    # def Select_Image(self,item):
    #     print(item)
    #     print(item.text())
    #     self.img_path =r''.join(item.text())
    #     self.image = cv2.imdecode(np.fromfile(self.img_path, dtype=np.uint8), -1)
    #     self.show_img([self.image])
    #     self.Start_Classify()
    def Row_Changed(self,row):
        print(row)
        self.img_path = r''.join(self.img_listWidget.item(row).text())
        self.image = cv2.imdecode(np.fromfile(self.img_path, dtype=np.uint8), -1)
        self.show_img([self.image])
        self.Start_Classify()
        # self.Start_Yolo()
        # self.Auto_Save_Xml()

    def Auto_Save_Xml(self):
        if self.T_AutoSave_checkBox.isChecked():
            img_name = os.path.basename(self.img_path)
            # save_image_path =
            w,h = self.image.shape[0],self.image.shape[1]
            write_xml(self.img_path,self.results,save_xml_folder=self.save_folder,width=w,height=h)
        pass

    def Save_Xml(self):
        w, h = self.image.shape[0], self.image.shape[1]
        basename = os.path.basename(self.img_path)
        write_xml(self.img_path, self.results, save_xml_folder=self.save_folder, width=w, height=h)

    def Auto_Run_AI(self):
        # count = self.
        self.T_AutoRunButton.setStyleSheet('background-color:green')
        self.progressBar.reset()
        save_txt_name = 'result.txt'
        f = open(os.path.join(self.img_folder,save_txt_name),'w',encoding='utf-8')
        f.write('image_path\timage_type\timage_result\n')
        count = self.img_listWidget.count()
        self.progressBar.setMaximum(count-1)

        for i in range(count):
            self.img_path = r''.join(self.img_listWidget.item(i).text())
            self.image = cv2.imdecode(np.fromfile(self.img_path, dtype=np.uint8), -1)
            self.show_img([self.image])
            self.Start_Classify()
            f.write('{}\t{}\t{}\n'.format(self.img_path,self.results['cls'][0],self.results['cls'][1]))
            self.progressBar.setValue(i)
            self.repaint()
        f.close()
        self.T_AutoRunButton.setStyleSheet(None)
    def clear_listwidget(self):
        try:
            self.img_listWidget.currentRowChanged.disconnect(self.Row_Changed)
            while (self.img_listWidget.count()>0):
                item = self.img_listWidget.takeItem(0)
                # del item
        except Exception as e:
            print(e)
